package com.sifou.courses.limit;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RedissonClient;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;

import javax.annotation.Resource;

/**
 * @author liuzhongxu
 * @date 2020/9/24
 */
@Slf4j
@Order(0)
@Aspect
@Component
public class RateLimitAspect {

    /**
     * single,cluster
     *  用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
     */
    @Resource
    private RateLimitMapping rateLimitMapping;

    @Resource
    private RedissonClient redissonClient;

    private long time;

    @Pointcut("@annotation(RateLimit)")
    public void aspect() {}

    @Around("aspect() && @annotation(limit)")
    public Object Interceptor(ProceedingJoinPoint proceedingJoinPoint, RateLimit limit) {

        MutiRateLimiter rateLimiter;
        Object result;
        String className = proceedingJoinPoint.getSignature().getDeclaringType().getSimpleName();
        String methodName = proceedingJoinPoint.getSignature().getName();
        String rateLimitName = String.format("%s.%s", className, methodName);

        if (RateLimit.LimitType.single.equals(limit.limitType())) {
            // 获取单机令牌桶
            rateLimiter = getRateLimiter(rateLimitName);
        } else if (RateLimit.LimitType.cluster.equals(limit.limitType())) {
            // 获取集群令牌桶
            rateLimiter = getRRateLimiter(rateLimitName);
        } else {
            log.warn("IllegalArgumentException. LimitType not defined!");
            return true;
        }

        if (rateLimiter.tryAcquire()) {
            log.debug("{} get 1 tokens: {}", rateLimitName, System.currentTimeMillis() - time);
            time = System.currentTimeMillis();
            try {
                result = proceedingJoinPoint.proceed();
            } catch (Throwable throwable) {
                throw new RuntimeException(throwable.getMessage());
            }
        } else {
            // 触发限流
            throw new RateLimitException("系统繁忙，请稍候重试~");
        }

        return result;
    }

    /**
     * 集群限流
     * @param rateLimitName
     * @return
     */
    private MutiRateLimiter getRRateLimiter(String rateLimitName) {
        RRateLimiter clusterRateLimiter = redissonClient.getRateLimiter(rateLimitName);
        return new MutiRateLimiter(clusterRateLimiter);
    }

    /**
     * 单机限流
     * @param rateLimitName
     * @return
     */
    private MutiRateLimiter getRateLimiter(String rateLimitName) {
        RateLimiter singleRateLimiter = rateLimitMapping.getSingleLimiterMap().get(rateLimitName);
        return new MutiRateLimiter(singleRateLimiter);
    }

    /**
     * 令牌桶
     */
    class MutiRateLimiter{
        private RateLimiter singleRateLimiter = null;
        private RRateLimiter clusterRateLimiter = null;

        MutiRateLimiter(RateLimiter singleRateLimiter) {
            this.singleRateLimiter = singleRateLimiter;
        }

        MutiRateLimiter(RRateLimiter clusterRateLimiter) {
            this.clusterRateLimiter = clusterRateLimiter;
        }

        /**
         * 获取令牌
         * @return
         */
        boolean tryAcquire() {
            if (!ObjectUtils.isEmpty(this.singleRateLimiter)) {
                return this.singleRateLimiter.tryAcquire();
            }
            if (!ObjectUtils.isEmpty(this.clusterRateLimiter)) {
                return this.clusterRateLimiter.tryAcquire();
            }
            return true;
        }
    }

}
